#include "armor_detector/ArmorDetectorNetwork.h"
namespace rm_auto_aim
{
ArmorDetectorNetwork::ArmorDetectorNetwork(int img_size_,std::string model_path_,int enemy_color_,float score_threshold_,float nms_threshold_):img_size_(img_size_),enemy_color_(enemy_color_),score_threshold_(score_threshold_),nms_threshold_(nms_threshold_) {
    compiled_model = core.compile_model(model_path_, "CPU"); // Compile the Model
    infer_request = compiled_model.create_infer_request();                            // Create an Inference Request

}
ArmorDetectorNetwork::~ArmorDetectorNetwork() {}

void ArmorDetectorNetwork::DrawResult(cv::Mat& __DebugNet_img,int id, float conf, std::vector<cv::Point2f> armor_verticles_) {
    std::vector<cv::Scalar> colors = {
        cv::Scalar(0, 0, 255), cv::Scalar(0, 255, 0), cv::Scalar(255, 0, 0), 
        cv::Scalar(255, 100, 50), cv::Scalar(50, 100, 255), cv::Scalar(255, 50, 100)
    };

    std::string label = class2name[id] + ":" + std::to_string(conf).substr(0, 4);
    cv::Size textSize = cv::getTextSize(label, cv::FONT_HERSHEY_SIMPLEX, 0.5, 1, 0);
    cv::Rect textBox(armor_verticles_[0].x, armor_verticles_[0].y - 15, textSize.width, textSize.height + 5);
    cv::rectangle(__DebugNet_img, textBox, colors[id % 6], cv::FILLED);
    putText(__DebugNet_img, label, cv::Point(armor_verticles_[0].x, armor_verticles_[0].y  - 5), cv::FONT_HERSHEY_SIMPLEX, 0.5, cv::Scalar(255, 255, 255));

    // keypoints
    for(int i = 0; i < 4; i++) {
        int x = armor_verticles_[i].x;
        int y = armor_verticles_[i].y;
        x = std::min(std::max(x, 0), __DebugNet_img.cols);
        y = std::min(std::max(y, 0), __DebugNet_img.rows);
        cv::circle(__DebugNet_img, cv::Point(x, y), 5, colors[i], -1);
        
        std::string point_label = std::to_string(i + 1);
        cv::putText(__DebugNet_img, point_label, cv::Point(x, y), cv::FONT_HERSHEY_SIMPLEX, 0.5, cv::Scalar(255, 255, 255));
    }
}

// Keep the ratio before resize
cv::Mat ArmorDetectorNetwork::letterbox(const cv::Mat &source) {
    int col = source.cols;
    int row = source.rows;
    int _max = MAX(col, row);
    cv::Mat result = cv::Mat::zeros(_max, _max, CV_8UC3);
    source.copyTo(result(cv::Rect(0, 0, col, row)));
    return result;
}

std::vector<Armor> ArmorDetectorNetwork::detect(cv::Mat& img) {
    // Preprocess the image
    cv::Mat letterbox_img = letterbox(img);
    float scale = letterbox_img.size[0] / (float)img_size_;
    cv::Mat blob = cv::dnn::blobFromImage(letterbox_img, 1.0 / 255.0, cv::Size(img_size_, img_size_), cv::Scalar(), true);
    ov::Output<const ov::Node> input_port = compiled_model.input();                              // Get input port for model with one input
    ov::Tensor input_tensor(input_port.get_element_type(), input_port.get_shape(), blob.ptr(0)); // Create tensor from external memory
    infer_request.set_input_tensor(input_tensor);                                                // Set input tensor for model with one input

    // Start inference
    infer_request.infer();
    // Get the inference result
    ov::Tensor output = infer_request.get_output_tensor(0);
    ov::Shape output_shape = output.get_shape();

    float *data = output.data<float>();
    cv::Mat output_buffer(output_shape[1], output_shape[2], CV_32F, data);
    transpose(output_buffer, output_buffer); //[8400,16]
    std::vector<int> class_ids;
    std::vector<float> class_scores;
    std::vector<cv::Rect> boxes;
    std::vector<std::vector<float>> objects_keypoints;

    // Figure out the bbox, class_id and class_score
    for(int i = 0; i < output_buffer.rows; i++) {
        cv::Mat classes_scores = output_buffer.row(i).colRange(4, 20);
        cv::Point class_id;
        double maxClassScore;
        minMaxLoc(classes_scores, 0, &maxClassScore, 0, &class_id);

        if(maxClassScore > score_threshold_) {
            class_scores.push_back(maxClassScore);
            class_ids.push_back(class_id.x);
            float cx = output_buffer.at<float>(i, 0);
            float cy = output_buffer.at<float>(i, 1);
            float w = output_buffer.at<float>(i, 2);
            float h = output_buffer.at<float>(i, 3);

            int left = int((cx - 0.5 * w) * scale);
            int top = int((cy - 0.5 * h) * scale);
            int width = int(w * scale);
            int height = int(h * scale);

            // Get the keypoints
            std::vector<float> keypoints;
            cv::Mat kpts = output_buffer.row(i).colRange(20, 28);
            for(int i = 0; i < 4; i++) {
                float x = kpts.at<float>(0, i * 2 + 0) * scale;
                float y = kpts.at<float>(0, i * 2 + 1) * scale;
                keypoints.push_back(x);
                keypoints.push_back(y);
            }

            boxes.push_back(cv::Rect(left, top, width, height));
            objects_keypoints.push_back(keypoints);
        }
    }

    // NMS
    std::vector<int> indices;
    cv::dnn::NMSBoxes(boxes, class_scores, score_threshold_, nms_threshold_, indices);
    std::vector<Armor> detected_armors_;
    // -------- Visualize the detection results -----------
    for(size_t i = 0; i < indices.size(); i++) {
        int index = indices[i];
        int class_id = class_ids[index];
        int color_idx = 1 - (class_id / 8);
        std::vector<float> object_keypoints = objects_keypoints[index];
        if(color_idx != enemy_color_) {
            continue;
        }
        Armor res;
        res.id_ = class_id % 8;

        int center_x = 0;
        int center_y = 0;
        for(int j = 0; j < 4; j++) {
            int x = (int)object_keypoints[j * 2];
            int y = (int)object_keypoints[j * 2 + 1];
            x = std::min(std::max(x, 0), img.cols);
            y = std::min(std::max(y, 0), img.rows);
            res.armor_verticles_[j].x = x;
            res.armor_verticles_[j].y = y;
            center_x += x / 4;
            center_y += y / 4;
        }
        float avg_light_length = \
            cv::norm(
                cv::Point2i(res.armor_verticles_[2].x, res.armor_verticles_[2].y) - \
                cv::Point2i(res.armor_verticles_[3].x, res.armor_verticles_[3].y) - \
                cv::Point2i(res.armor_verticles_[0].x, res.armor_verticles_[0].y) + \
                cv::Point2i(res.armor_verticles_[1].x, res.armor_verticles_[1].y)
            ) / 2;
        float light_width = \
            cv::norm(
                res.armor_verticles_[3] + res.armor_verticles_[2] - \
                res.armor_verticles_[1] - res.armor_verticles_[0]
            ) / 2;
        if(light_width / avg_light_length > 3.2) {
            res.type = ArmorType::LARGE;
        } else {
            res.type = ArmorType::SMALL;
        }
        detected_armors_.emplace_back(res);
        
    }



    return detected_armors_;
}
}